#pragma once
#include <vector>
#include <numeric>
#include <cuda_runtime.h> // cuda相关头文件
#include "NvInfer.h"
#include "dlm_crypt.hpp"
#include "cuda_utility.hpp"
#include "logging.h"

inline int64_t static volume(const nvinfer1::Dims& d)
{
	return std::accumulate(d.d, d.d + d.nbDims, 1, std::multiplies<int64_t>());
}
inline unsigned int static getElementSize(nvinfer1::DataType t)
{
	switch (t)
	{
	case nvinfer1::DataType::kINT32: return 4;
	case nvinfer1::DataType::kFLOAT: return 4;
	case nvinfer1::DataType::kHALF: return 2;
	case nvinfer1::DataType::kBOOL:
	case nvinfer1::DataType::kINT8: return 1;
	}
	throw std::runtime_error("Invalid DataType.");
	return 0;
}

class InferBase
{
public:
	InferBase() {
		cudaStreamCreate(&m_stream);
	}
	virtual	~InferBase() {
		cudaStreamDestroy(m_stream);
		Reset();
	}
	virtual bool Init(int gpu_id, const char* dats_path)
	{
		std::string file_path = dats_path;
		std::vector<char> v_engine_data;
		if (!set_gpu_id(gpu_id) || !decryptFile(file_path, m_info, v_engine_data))
		{
			return false;
		}
		return LoadEngine(v_engine_data);
	}
	virtual int BatchSize()
	{
		return m_info.batch_size;
	}
	virtual void Reset()
	{
		if (m_context)
		{
			m_context->destroy();
		}
		if (m_engine)
		{
			m_engine->destroy();
		}
		for (int i = 0; i < m_v_dev_buffer.size(); ++i)
		{
			if (m_v_dev_buffer[i])
			{
				cudaFree(m_v_dev_buffer[i]);
				m_v_dev_buffer[i] = 0;
			}
		}
		m_v_dev_buffer.clear();
		m_v_dev_buf_size.clear();
	}
	virtual bool LoadEngine(std::vector<char> v_engine_data)
	{
		Reset();
		nvinfer1::IRuntime* trtRuntime;
		trtRuntime = nvinfer1::createInferRuntime(sample::gLogger.getTRTLogger());

		m_engine = trtRuntime->deserializeCudaEngine(v_engine_data.data(), v_engine_data.size(), nullptr);
		if (m_engine == nullptr)
		{
			std::cout << "deserialize fail" << std::endl;
			return false;
		}
		m_context = m_engine->createExecutionContext();
		if (m_context == nullptr)
		{
			std::cout << "create context fail" << std::endl;
			return false;
		}
		return true;
	}

protected:
	st_encrypt_info m_info;

	cudaStream_t m_stream;
	nvinfer1::ICudaEngine* m_engine = nullptr;
	nvinfer1::IExecutionContext* m_context = nullptr;

	std::vector<void*> m_v_dev_buffer;
	std::vector<int64_t> m_v_dev_buf_size;
private:
};